/*
 * QTSM Device Driver Implement
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 *
 * Lots of code copied from AWS nsm driver, and modified to be useful in qingtian_enclaves:
 * https://github.com/aws/aws-nitro-enclaves-sdk-bootstrap/blob/main/nsm-driver/nsm.c
 * Amazon Nitro Secure Module driver.
 * Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 */
#include <linux/file.h>
#include <linux/fs.h>
#include <linux/interrupt.h>
#include <linux/miscdevice.h>
#include <linux/module.h>
#include <linux/mutex.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/uaccess.h>
#include <linux/uio.h>
#include <linux/virtio.h>
#include <linux/virtio_config.h>
#include <linux/sched.h>
#include <linux/wait.h>
#include <linux/pm.h>
#include "qtsm.h"

#define VIRTIO_ID_QTSM                  63

/* Register this as a misc driver */
#define QTSM_DEV_MINOR                  3
#define QTSM_DEV_NAME                   "qtsm"
#define QTSM_DEV_DATA_OFFSET            0x100
#define QTSM_IOCTL_MAGIC                0xC7
#define QTSM_IO_REQUEST                 _IOWR(QTSM_IOCTL_MAGIC, 0, struct qtsm_message)
#define QTSM_REQUEST_MAX_SIZE           0x1000
#define QTSM_RESPONSE_MAX_SIZE          0x8000

/* Timeout for QTSM virtqueue response in milliseconds. */
#define QTSM_DEFAULT_TIMEOUT_MSECS      120000 /* 2 minutes */

/* Virtio MMIO device */
struct virtio_qtsm {
	struct virtio_device *vdev;
	struct virtqueue *vq;
};

/* The name of the QTSM device virtqueue */
const char *QTSM_VQ_NAME = "qtsm.vq.0";

/* QTSM device ID */
static const struct virtio_device_id id_table[] = {
	{ VIRTIO_ID_QTSM, VIRTIO_DEV_ANY_ID },
	{ 0 },
};

static struct virtio_qtsm *qtsm_dev;
static struct mutex qtsm_lock;
static wait_queue_head_t qtsm_waitqueue;
static bool qtsm_device_notified;

/* QTSM message from user-space */
struct qtsm_message {
	/* Request from user */
	struct iovec request;
	/* Response to user */
	struct iovec response;
};

/* Copy of QTSM message in kernel-space */
struct qtsm_kernel_message {
	/* Copy of user request in kernel memory */
	struct kvec request;
	/* Copy of user response in kernel memory */
	struct kvec response;
};

/* Get the virtqueue */
static inline struct virtqueue *qtsm_get_vq(struct virtio_qtsm *qtsm)
{
	return qtsm->vq;
}

/* Virtqueue interrupt handler */
static void qtsm_vq_callback(struct virtqueue *vq)
{
	pr_info("qtsm_vq_callback start.\n");
	qtsm_device_notified = true;
	wake_up(&qtsm_waitqueue);
}

static int qtsm_device_init_vq(struct virtio_qtsm *qtsm)
{
	struct virtqueue *vq;

	vq = virtio_find_single_vq(qtsm->vdev, qtsm_vq_callback, QTSM_VQ_NAME);
	if (IS_ERR(vq))
		return PTR_ERR(vq);

	qtsm->vq = vq;

	return 0;
}

/* Handler for probing the QTSM device */
static int qtsm_dev_probe(struct virtio_device *vdev)
{
	struct virtio_qtsm *qtsm = NULL;
	int rc;

	vdev->priv = qtsm = kzalloc(sizeof(*qtsm), GFP_KERNEL);
	if (!qtsm)
		return -ENOMEM;

	qtsm->vdev = vdev;

	/* register the virtqueue */
	rc = qtsm_device_init_vq(qtsm);
	if (rc) {
		vdev->priv = NULL;
		kfree(qtsm);
		return rc;
	}

	qtsm_dev = qtsm;
	pr_info("QTSM device has been probed.\n");
	return 0;
}

/* Handler for removing the QTSM device */
static void qtsm_dev_remove(struct virtio_device *vdev)
{
	struct virtio_qtsm *qtsm = vdev->priv;

	vdev->config->del_vqs(vdev);
	kfree(qtsm);
	qtsm_dev = NULL;
	pr_info("QTSM device has been removed.\n");
}

/* Handler for changing the configuration of the QTSM device */
static void qtsm_dev_cfg_changed(struct virtio_device *vdev)
{
	pr_info("QTSM device configuration has been changed.\n");
}

static int qtsm_file_open(struct inode *node, struct file *file)
{
	pr_info("QTSM device file opened.\n");
	return 0;
}

static int qtsm_file_close(struct inode *inode, struct file *file)
{
	pr_info("QTSM device file closed.\n");
	return 0;
}

/* Copy an entire message from user-space to kernel-space */
static int message_copy_from_user(struct qtsm_kernel_message *dst, struct qtsm_message *src)
{
	struct qtsm_message shallow_copy;

	if (!src || !dst)
		return -EINVAL;

	/* The destination's request and response buffers should be NULL. */
	if (dst->request.iov_base || dst->response.iov_base)
		return -EINVAL;

	/* First, make a shallow copy to be able to read the inner pointers */
	if (copy_from_user(&shallow_copy, src, sizeof(shallow_copy)) != 0)
		return -EINVAL;

	/* Verify the user input size. */
	if (shallow_copy.request.iov_len > QTSM_REQUEST_MAX_SIZE)
		return -EMSGSIZE;

	/* Allocate kernel memory for the user request */
	dst->request.iov_len = shallow_copy.request.iov_len;
	dst->request.iov_base = kmalloc(dst->request.iov_len, GFP_KERNEL);
	if (!dst->request.iov_base)
		return -ENOMEM;

	/* Copy the request content */
	if (copy_from_user(dst->request.iov_base, shallow_copy.request.iov_base,
					   dst->request.iov_len) != 0) {
		kfree(dst->request.iov_base);
		return -EFAULT;
	}

	/* Allocate kernel memory for the response, up to a fixed limit */
	dst->response.iov_len = shallow_copy.response.iov_len;
	if (dst->response.iov_len > QTSM_RESPONSE_MAX_SIZE)
		dst->response.iov_len = QTSM_RESPONSE_MAX_SIZE;

	dst->response.iov_base = kmalloc(dst->response.iov_len, GFP_KERNEL);
	if (!dst->response.iov_base) {
		kfree(dst->request.iov_base);
		return -ENOMEM;
	}

	return 0;
}

/* Copy a message back to user-space */
static int message_copy_to_user(struct qtsm_message *user_msg,
	struct qtsm_kernel_message *kern_msg)
{
	struct qtsm_message shallow_copy;

	if (!kern_msg || !user_msg)
		return -EINVAL;

	/*
	* First, do a shallow copy of the user-space message. This is needed in
	* order to get the request block data, which we do not need to copy but
	* must preserve in the message sent back to user-space.
	*/
	if (copy_from_user(&shallow_copy, user_msg, sizeof(shallow_copy)) != 0)
		return -EINVAL;

	/* Do not exceed the capacity of the user-provided response buffer */
	shallow_copy.response.iov_len = kern_msg->response.iov_len;

	/* Only the response content must be copied back to user-space */
	if (copy_to_user(shallow_copy.response.iov_base,
					 kern_msg->response.iov_base, shallow_copy.response.iov_len) != 0)
		return -EINVAL;

	if (copy_to_user(user_msg, &shallow_copy, sizeof(shallow_copy)) != 0)
		return -EFAULT;

	return 0;
}

/* Delete a message stored in kernel-space */
static void message_delete(struct qtsm_kernel_message *message)
{
	if (!message)
		return;

	kfree(message->request.iov_base);
	kfree(message->response.iov_base);
	message->request.iov_base = message->response.iov_base = NULL;
}

static int check_communicate_with_device_input(struct qtsm_kernel_message *message, struct virtqueue **vq)
{
	if (!message) {
		pr_err("qtsm_communicate_with_device: message is NULL.\n");
		return -EINVAL;
	}

	if (!qtsm_dev)
		return -ENXIO;

	*vq = qtsm_get_vq(qtsm_dev);
	if (!(*vq))
		return -ENXIO;

	/* Verify if buffer memory is valid. */
	if (!virt_addr_valid(message->request.iov_base) ||
		!virt_addr_valid(((u8 *)message->request.iov_base) + message->request.iov_len - 1) ||
		!virt_addr_valid(message->response.iov_base) ||
		!virt_addr_valid(((u8 *)message->response.iov_base) + message->response.iov_len - 1))
		return -EINVAL;
	return 0;
}

static void qtsm_clean_virtqueue(struct virtqueue *vq, unsigned int *len)
{
	while (virtqueue_get_buf(vq, len) != NULL);
	pr_info("already cleanup.\n");
}

/* Forward a message to the QTSM device and wait for the response from it */
static int qtsm_communicate_with_device(struct qtsm_kernel_message *message, unsigned int cmd)
{
	struct virtqueue *vq = NULL;
	struct scatterlist sg_in, sg_out;
	unsigned int len;
	void *queue_buf = NULL;
	bool kicked = false;
	int rc;

	rc = check_communicate_with_device_input(message, &vq);
	if (rc)
		return rc;

	/* Initialize scatter-gather lists with request and response buffers. */
	sg_init_one(&sg_out, message->request.iov_base, message->request.iov_len);
	sg_init_one(&sg_in, message->response.iov_base, message->response.iov_len);

	/* Add the request buffer (read by the device). */
	rc = virtqueue_add_outbuf(vq, &sg_out, 1, message->request.iov_base, GFP_KERNEL);
	if (rc)
		return rc;

	/* Add the response buffer (written by the device). */
	rc = virtqueue_add_inbuf(vq, &sg_in, 1, message->response.iov_base, GFP_KERNEL);
	if (rc)
		goto cleanup;

	qtsm_device_notified = false;
	kicked = virtqueue_kick(vq);
	if (!kicked) {
		/* Cannot kick the virtqueue. */
		pr_err("failed to kick the virtqueue.\n");
		rc = -EIO;
		goto cleanup;
	}

	/* If the kick succeeded, wait for the device's response. */
	rc = wait_event_timeout(qtsm_waitqueue, qtsm_device_notified == true, msecs_to_jiffies(QTSM_DEFAULT_TIMEOUT_MSECS));
	if (!rc) {
		pr_err("failed to wait reply(timeout).\n");
		rc = -ETIMEDOUT;
		goto cleanup;
	}

	queue_buf = virtqueue_get_buf(vq, &len);
	if (!queue_buf || (queue_buf != message->request.iov_base)) {
		pr_err("QTSM device received wrong request buffer.\n");
		rc = -ENODATA;
		goto cleanup;
	}

	queue_buf = virtqueue_get_buf(vq, &len);
	if (!queue_buf || (queue_buf != message->response.iov_base)) {
		pr_err("QTSM device received wrong response buffer.\n");
		rc = -ENODATA;
		goto cleanup;
	}

	/* Make sure the response length doesn't exceed the buffer capacity. */
	if (len < message->response.iov_len)
		message->response.iov_len = len;

	return 0;

cleanup:
	qtsm_clean_virtqueue(vq, &len);
	return rc;
}

static long qtsm_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
	struct qtsm_kernel_message message;
	int status = 0;

	if (!qtsm_dev) {
		return -EINVAL;
	}

	if (cmd != QTSM_IO_REQUEST)
		return -EINVAL;

	/* The kernel message structure must be cleared */
	memset(&message, 0, sizeof(message));

	/* Copy the message from user-space to kernel-space */
	status = message_copy_from_user(&message, (struct qtsm_message *)arg);
	if (status != 0)
		return status;

	/* Communicate with the QTSM device */
	mutex_lock(&qtsm_lock);
	status = qtsm_communicate_with_device(&message, QTSM_IO_REQUEST);
	mutex_unlock(&qtsm_lock);

	if (status != 0) {
		pr_err("qtsm_communicate_with_device failed!\n");
		message_delete(&message);
		return status;
	}

	/* Copy the response back to user-space */
	status = message_copy_to_user((struct qtsm_message *)arg, &message);

	/* At this point, everything succeeded, so clean up and finish. */
	message_delete(&message);

	return status;
}

/* Virtio QTSM device configuration structure */
static struct virtio_driver virtio_qtsm_driver = {
	.feature_table              = 0,
	.feature_table_size         = 0,
	.feature_table_legacy       = 0,
	.feature_table_size_legacy  = 0,
	.driver.name                = KBUILD_MODNAME,
	.driver.owner               = THIS_MODULE,
	.id_table                   = id_table,
	.probe                      = qtsm_dev_probe,
	.remove                     = qtsm_dev_remove,
	.config_changed             = qtsm_dev_cfg_changed
};

/* Supported driver operations. */
static const struct file_operations qtsm_fops = {
	.open           = qtsm_file_open,
	.release        = qtsm_file_close,
	.unlocked_ioctl = qtsm_ioctl,
};

/* Driver configuration. */
static struct miscdevice qtsm_misc_dev = {
	.minor  = QTSM_DEV_MINOR,
	.name   = QTSM_DEV_NAME,
	.fops   = &qtsm_fops,
	.mode   = 0666
};

void qtsm_power_off(void)
{
	__u32 cmd = VIRTIO_QTSM_CONFIG_CMD_POWER_OFF;

	if (unlikely(!qtsm_dev)) {
		pr_err("Power off: qtsm device error.\n");
		return;
	}

	pr_info("QTSM: power off requested.\n");
	virtio_cwrite(qtsm_dev->vdev, struct virtio_qtsm_config,
		      cmd, &cmd);
}

static int __init qtsm_driver_init(void)
{
	int rc;

	mutex_init(&qtsm_lock);
	init_waitqueue_head(&qtsm_waitqueue);

	rc = misc_register(&qtsm_misc_dev);
	if (rc)
		pr_err("QTSM driver initialization error: %d.\n", rc);
	else
		pr_info("QTSM driver version = 1.0.%d.\n", QTSM_DEV_MINOR);

	qtsm_dev = NULL;
	rc = register_virtio_driver(&virtio_qtsm_driver);
	if (rc) {
		pr_err("QTSM device initialization error: %d.\n", rc);
		misc_deregister(&qtsm_misc_dev);
	} else {
		pr_info("QTSM device ID = 0x%X.\n", virtio_qtsm_driver.id_table[0].device);
	}
	pm_power_off = qtsm_power_off;
	pr_info("QTSM set power off callback.\n");

	return rc;
}

static void __exit qtsm_driver_exit(void)
{
	unregister_virtio_driver(&virtio_qtsm_driver);
	misc_deregister(&qtsm_misc_dev);
	mutex_destroy(&qtsm_lock);
	pr_info("QTSM driver exited.\n");
}

module_init(qtsm_driver_init);
module_exit(qtsm_driver_exit);
MODULE_DEVICE_TABLE(virtio, id_table);
MODULE_DESCRIPTION("Virtio qtsm driver");
MODULE_LICENSE("GPL");
